# #!/usr/bin/env python3
# import logging
# logging.basicConfig(level=logging.DEBUG)

# from pydantic import BaseModel

# from outlines import models, generate


# class User(BaseModel):
#     name: str
#     last_name: str
#     id: int
#     location: str


model_path = "/home/wangxianda/model_zoo/llama_3.2_3b_instruct"
# model = models.transformers(model_path)
# generator = generate.json(model, User)
# result = generator(
#     "Create a user profile with the fields name, last_name and id, his name is David"
# )
# print(result)
# User(name="John", last_name="Doe", id=11)
from pydantic import BaseModel
import torch

from outlines import models
from outlines import generate

model = models.transformers(model_path,
                            device="auto",
                            model_kwargs={"torch_dtype": torch.float16})

schema = """
{
  "title": "Tool",
  "type": "object",
  "properties": {
    "name": {
      "type": "string",
      "enum": ["Chat", "Calculator", "Ping"]
    },
    "action_input": {
      "type": "string"
    }
  },
  "required": ["name", "action_input"]
}
"""

generator = generate.json(model, schema)
result = generator(
    "choose a network status json, ping 192.168.3.2", max_tokens=100
)
print(result)
# User(name="John", last_name="Doe", id=11)

################
################
################
# from outlines import models
# from outlines import generate

# def add(a: int, b: int):
#     return a + b

# model = models.transformers("microsoft/Phi-3-mini-4k-instruct")
# generator = generate.json(model, add)
# result = generator("Return two integers named a and b respectively. a is odd and b even.")

# print(add(**result))
# 3
